/**
 * 
 */
package edu.arizona.bio5.onto.decomposition.persistent.rdf;


import java.io.Reader;
import java.io.StringReader;
import java.io.StringWriter;
import java.net.URI;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.openrdf.model.Value;
import org.openrdf.model.ValueFactory;
import org.openrdf.model.vocabulary.RDF;
import org.openrdf.query.BindingSet;
import org.openrdf.query.GraphQuery;
import org.openrdf.query.MalformedQueryException;
import org.openrdf.query.Query;
import org.openrdf.query.QueryEvaluationException;
import org.openrdf.query.QueryLanguage;
import org.openrdf.query.TupleQuery;
import org.openrdf.query.TupleQueryResult;
import org.openrdf.repository.Repository;
import org.openrdf.repository.RepositoryConnection;
import org.openrdf.repository.RepositoryException;
import org.openrdf.rio.RDFHandler;
import org.openrdf.rio.RDFHandlerException;
import org.openrdf.rio.n3.N3Writer;
import org.semanticweb.owlapi.apibinding.OWLManager;
import org.semanticweb.owlapi.io.ReaderDocumentSource;
import org.semanticweb.owlapi.model.OWLAxiom;
import org.semanticweb.owlapi.model.OWLDataFactory;
import org.semanticweb.owlapi.model.OWLEntity;
import org.semanticweb.owlapi.model.OWLOntology;
import org.semanticweb.owlapi.model.OWLOntologyCreationException;
import org.semanticweb.owlapi.model.OWLOntologyManager;

import uk.ac.manchester.cs.graph.DirectedGraph;
import uk.ac.manchester.cs.graph.DirectedGraphImpl;
import uk.ac.manchester.cs.owlapi.modularity.ModuleType;
import edu.arizona.bio5.onto.decomposition.Atom;
import edu.arizona.bio5.onto.decomposition.MSSLabel;
import edu.arizona.bio5.onto.decomposition.persistent.ADInfo;
import edu.arizona.bio5.onto.decomposition.persistent.ADLoader;
import edu.arizona.bio5.onto.decomposition.persistent.ADPersistenceException;
import edu.arizona.bio5.onto.decomposition.persistent.AbstractAtomLoader;
import edu.arizona.bio5.onto.decomposition.persistent.PersistentAtom;
import edu.arizona.bio5.onto.decomposition.persistent.PersistentAtomicDecomposition;

/**
 * Loads atomic decomposition from an RDF database through Sesame's RepositoryConnection
 * 
 * @author Pavel Klinov
 *
 */
public class SesameADLoader implements ADLoader {

	static final String PREFIX = 	"PREFIX ad: <" + RDFADConstants.AD_BASE_NS + "> \n" + 
											"PREFIX ads: <" + RDFADConstants.AD_BASE_SCHEMA + "> \n" +
											"PREFIX rdf: <" + RDF.NAMESPACE + "> \n\n";
	static final String AD_INFO_QUERY = PREFIX + "SELECT ?mt FROM \n" +
															"<ng> \n" + 
															"WHERE {\n " +
															"?onto_uri ad:has-module-type ?mt \n" +
															"}";
	static final String AD_GRAPH_QUERY = PREFIX + 
			"SELECT ?atom ?dep FROM \n" +
			"<ng> \n" + 
			"WHERE {\n " +
			"?atom rdf:type ads:Atom \n" +
			"OPTIONAL" +
			"{ ?atom ads:dependsOn ?dep }" +
			"}";
	
	static final String AD_MSS_LABELS_QUERY = PREFIX + 
			"SELECT ?atom ?mss ?term FROM \n" +
			"<ng> \n" + 
			"WHERE {\n " +
			"?atom rdf:type ads:Atom .\n" +
			"?atom ads:hasLabel ?label .\n" +
			"?label rdf:type ads:MinSeedSignatureLabel .\n" +
			"?label ads:hasSignature ?mss .\n" +
			"?mss ads:hasTerm ?term .\n" +
			"} ORDER BY ?atom ?mss";
	
	static final String AD_TERM_INDEX_QUERY = PREFIX + 
			"SELECT ?term ?atom FROM \n" +
			"<ng> \n" + 
			"WHERE {\n " +
			"?term ads:hasRelevantAtoms ?container.\n" +
			"?container ads:containsAtom ?atom .\n" +
			"} ORDER BY ?term";	
	
	static final String AD_ATOM_GRAPH_QUERY = PREFIX + 
			"CONSTRUCT { ?s ?p ?o } FROM \n" +
			"<ng> \n" + 
			"WHERE {\n " +
			"?s ?p ?o \n" +
			"}";	
	
	private final Repository mRepo;
	private RepositoryConnection mConnection;
	private final ValueFactory mFactory;
	private final String mOntoURI;
	
	public SesameADLoader(final Repository theRepo, final String ontoURI) {
		mRepo = theRepo;
		mFactory = mRepo.getValueFactory();
		mOntoURI = ontoURI;
	}
	
	private Query getQuery(String queryString, String ontoURI, boolean graph) throws RepositoryException, MalformedQueryException {
		String qString = queryString.replaceAll("<ng>", "<" + RDFADConstants.AD_BASE_NS + ontoURI + ">");
		
		return graph ? mConnection.prepareGraphQuery(QueryLanguage.SPARQL, qString) : mConnection.prepareTupleQuery(QueryLanguage.SPARQL, qString);
	}
	
	@Override
	public PersistentAtomicDecomposition load() throws ADPersistenceException {
		
		try {
			mConnection = mRepo.getConnection();
			
			System.out.println("Loading AD info");
			
			ADInfo adInfo = loadADInfo();
			SesameAtomicDecompositionImpl ad = new SesameAtomicDecompositionImpl(mRepo, URI.create(mOntoURI), adInfo.getModType());			
			Map<Object, Atom> idAtomMap = new HashMap<Object, Atom>();
			
			System.out.println("Loading the AD graph");
			
			DirectedGraph<Atom, Object> adGraph = loadGraph(idAtomMap, ad);
			
			System.out.println("Loading labels");
			
			loadMSSLabels(idAtomMap);
			
			System.out.println("Loading the term based index");
			
			Map<URI, Set<Atom>> tbIndex = loadTermBasedIndex(idAtomMap);
			
			System.out.println("Loading global axioms");
			
			Set<OWLAxiom> globalAxioms = loadGlobalAxioms();
		
			ad.setADGraph(adGraph);
			ad.setTermBasedIndex(tbIndex);
			ad.setGlobalAxioms(globalAxioms);
			
			return ad;
			
		} catch (RepositoryException e) {
			throw new ADPersistenceException(e);
		} catch (QueryEvaluationException e) {
			throw new ADPersistenceException(e);
		}
		finally {
			closeConnection();
		}
	}
	
	private void closeConnection() {
		try {
			if (mConnection != null) {
				mConnection.close();
			}
		} catch (RepositoryException e) {}
	}
	

	private Map<URI, Set<Atom>> loadTermBasedIndex(Map<Object, Atom> idAtomMap) throws RepositoryException, QueryEvaluationException {
		try {
			Map<URI, Set<Atom>> tbIndex = new HashMap<URI, Set<Atom>>();
			TupleQuery query = (TupleQuery)getQuery(AD_TERM_INDEX_QUERY, mOntoURI, false);
			TupleQueryResult results = query.evaluate();
			Value currTerm = null;
			Set<Atom> termAtoms = new HashSet<Atom>();

			while (results.hasNext()) {
				BindingSet bindings = results.next();
				Value atomURI = bindings.getValue("atom");
				Value termURI = bindings.getValue("term");
				
				if (currTerm == null) currTerm = termURI;
				
				if (!currTerm.stringValue().equals(termURI.stringValue())) {
					//add relevant atoms
					tbIndex.put(java.net.URI.create(currTerm.stringValue()), termAtoms);
					termAtoms = new HashSet<Atom>();
					currTerm = termURI;
				}

				//add a new atom to the current set of relevant atoms
				termAtoms.add(idAtomMap.get(getAtomID(atomURI)));

			}
			//the last term
			tbIndex.put(java.net.URI.create(currTerm.stringValue()), termAtoms);

			return tbIndex;
			
		} catch (MalformedQueryException e) {
			throw new RuntimeException(e);
		}
	}

	private Set<OWLAxiom> loadGlobalAxioms() throws ADPersistenceException {

		try {
			String context = RDFADConstants.AD_BASE_NS + mOntoURI + "/" + RDFADConstants.AD_GLOBAL_AXIOMS_GRAPH;
			String qString = AD_ATOM_GRAPH_QUERY.replaceAll("<ng>", "<" + context + ">");				
			GraphQuery query = mConnection.prepareGraphQuery(QueryLanguage.SPARQL, qString);
			StringWriter strWriter = new StringWriter();
			RDFHandler writer = new N3Writer(strWriter);

			query.evaluate(writer);
			strWriter.flush();
			
			StringBuffer buffer = strWriter.getBuffer();
			
			OWLOntologyManager mgr = OWLManager.createOWLOntologyManager();
			OWLOntology atomOntology = mgr.loadOntologyFromOntologyDocument(new ReaderDocumentSource(new StringReader(buffer.toString())));
				
			return atomOntology.getAxioms();

		} catch (RepositoryException e) {
			throw new ADPersistenceException(e);
		} catch (MalformedQueryException e) {
			throw new RuntimeException(e);
		} catch (QueryEvaluationException e) {
			throw new ADPersistenceException(e);
		} catch (RDFHandlerException e) {
			throw new ADPersistenceException(e);
		} catch (OWLOntologyCreationException e) {
			throw new ADPersistenceException(e);
		}
	}

	private void loadMSSLabels(Map<Object, Atom> idAtomMap) throws RepositoryException, QueryEvaluationException, ADPersistenceException {
		try {
			OWLDataFactory factory = OWLManager.getOWLDataFactory();
			TupleQuery query = (TupleQuery)getQuery(AD_MSS_LABELS_QUERY, mOntoURI, false);

			TupleQueryResult results = query.evaluate();
			Value currAtom = null;
			Value currMSS = null;
			List<Set<OWLEntity>> mssList = new ArrayList<Set<OWLEntity>>();
			Set<OWLEntity> mss = new HashSet<OWLEntity>();

			while (results.hasNext()) {
				BindingSet bindings = results.next();
				Value atomURI = bindings.getValue("atom");
				Value mssURI = bindings.getValue("mss");
				Value termURI = bindings.getValue("term");
				
				if (currAtom == null) currAtom = atomURI;
				if (currMSS == null) currMSS = mssURI;
				
				if (!atomURI.stringValue().equals(currAtom.stringValue())) {
					//set the label
					mssList.add(mss);
					setLabel(currAtom, mssList, idAtomMap);
					mssList = new ArrayList<Set<OWLEntity>>();
					mss = new HashSet<OWLEntity>();
					currAtom = atomURI;
					currMSS = mssURI;
				}
				
				if (!currMSS.stringValue().equals(mssURI.stringValue())) {
					//start a new MSS
					mssList.add(mss);
					mss = new HashSet<OWLEntity>();
					currMSS = mssURI;
				}
				//add a new term to the current MSS
				mss.add(SesameOWLRDFUtils.termFromURI(factory, termURI.stringValue()));

			}
			//the last label
			mssList.add(mss);
			setLabel(currAtom, mssList, idAtomMap);
			
		} catch (MalformedQueryException e) {
			throw new RuntimeException(e);
		}
	}

	private void setLabel(Value currAtom, List<Set<OWLEntity>> mssList, Map<Object, Atom> idAtomMap) throws ADPersistenceException {
		Atom atom = idAtomMap.get(getAtomID(currAtom));
		
		if (atom != null) {
			atom.setLabel(new MSSLabel(mssList));
		}
		else {
			throw new ADPersistenceException("Couldn't find the label for " + currAtom.stringValue());
		}
	}

	private DirectedGraph<Atom, Object> loadGraph(Map<Object, Atom> idAtomMap, SesameAtomicDecompositionImpl ad) throws RepositoryException, QueryEvaluationException {
		try {
			TupleQuery query = (TupleQuery)getQuery(AD_GRAPH_QUERY, mOntoURI, false);
			DirectedGraph<Atom, Object> graph = new DirectedGraphImpl<Atom, Object>();
			
			query.setBinding("onto_uri", mFactory.createURI(RDFADConstants.AD_BASE_NS + mOntoURI));
			
			TupleQueryResult results = query.evaluate();
			
			while (results.hasNext()) {
				BindingSet bindings = results.next();
				//according the query we either get an atom declaration or a dependency statement
				Value atom = bindings.getValue("atom");
				Value dep = bindings.getValue("dep");
				//create or find the atom
				Atom graphAtom = findOrCreateAtom(atom, idAtomMap, ad);
				
				graph.addVertex(graphAtom);
				
				if (dep != null) {
					Atom graphDepAtom = findOrCreateAtom(dep, idAtomMap, ad);
					
					graph.addVertex(graphDepAtom);
					graph.addEdge("", graphAtom, graphDepAtom);
				}
				
				//System.out.println(bindings.getBinding("atom") + ", " + bindings.getBinding("dep"));
			}
			
			return graph;
		} catch (MalformedQueryException e) {
			throw new RuntimeException(e);
		}
	}

	private Atom findOrCreateAtom(Value atom, Map<Object, Atom> idAtomMap, SesameAtomicDecompositionImpl ad) {
		//get the id
		Object id = getAtomID(atom);
		Atom graphAtom = idAtomMap.get(id);
		
		graphAtom = graphAtom == null 
				? new PersistentAtom(id, new SesameAtomLoader(ad)) 
				: graphAtom;
		idAtomMap.put(id, graphAtom);
		
		return graphAtom;
	}
	
	static Object getAtomID(Value atomURI) {
		return atomURI.stringValue().substring(RDFADConstants.AD_ATOM_INSTANCE.length());
	}

	private ADInfo loadADInfo() throws RepositoryException, QueryEvaluationException, ADPersistenceException {
		try {
			TupleQuery query = (TupleQuery)getQuery(AD_INFO_QUERY, mOntoURI, false);
			
			query.setBinding("onto_uri", mFactory.createURI(RDFADConstants.AD_BASE_NS + mOntoURI));
			
			TupleQueryResult results = query.evaluate();
			Value modType = null;
			
			while (modType == null && results.hasNext()) {
				BindingSet bindings = results.next();
				
				modType = bindings.getValue("mt");
			}
			
			if (modType != null) {
				return new ADInfo(mOntoURI, ModuleType.valueOf(modType.stringValue()));
			}
			else {
				throw new ADPersistenceException("Failed to load the module type for AD");
			}
			
		} catch (MalformedQueryException e) {
			throw new RuntimeException(e);
		}
	}
}

class SesameAtomLoader extends AbstractAtomLoader {
	
	private final SesameAtomicDecompositionImpl mAD;
	
	SesameAtomLoader(SesameAtomicDecompositionImpl ad) {
		mAD = ad;
	}
	
	@Override
	protected Reader getOntologySourceReader(PersistentAtom atom) throws ADPersistenceException {
		// Run a graph query, serialize all statements into a string and
		// parse it back via the OWL API
		// TODO Find a way w/o serializing into RDF/XML
		try {
			RepositoryConnection conn = mAD.getConnection();
			String qString = SesameADLoader.AD_ATOM_GRAPH_QUERY.replaceAll("<ng>", "<" + SesameOWLRDFUtils.getAtomURI(atom, conn.getValueFactory()) + ">");				
			GraphQuery atomQuery = conn.prepareGraphQuery(QueryLanguage.SPARQL, qString);
			
			StringWriter strWriter = new StringWriter();
			RDFHandler writer = new N3Writer(strWriter);

			atomQuery.evaluate(writer);
			strWriter.flush();
			
			StringBuffer buffer = strWriter.getBuffer();
			
			return new StringReader(buffer.toString());

		} catch (RepositoryException e) {
			throw new ADPersistenceException(e);
		} catch (MalformedQueryException e) {
			throw new RuntimeException(e);
		} catch (QueryEvaluationException e) {
			throw new ADPersistenceException(e);
		} catch (RDFHandlerException e) {
			throw new ADPersistenceException(e);
		}
	}

/*	@Override
	public void load(PersistentAtom atom) throws ADPersistenceException {
		try {
			OWLOntologyManager mgr = OWLManager.createOWLOntologyManager();
			RepositoryConnection conn = mAD.getConnection();
			String qString = SesameADLoader.AD_ATOM_GRAPH_QUERY.replaceAll("<ng>", "<" + SesameRDFUtils.getAtomURI(atom, conn.getValueFactory()) + ">");				
			GraphQuery atomQuery = conn.prepareGraphQuery(QueryLanguage.SPARQL, qString);
			OWLOntology atomOntology = mgr.createOntology();
			OWLRDFConsumer consumer = new OWLRDFConsumer(atomOntology, new AnonymousNodeChecker() {
				
				@Override
				public boolean isAnonymousSharedNode(String arg) {
					return false;
				}
				
				@Override
				public boolean isAnonymousNode(String arg) {
					return false;
				}
				
				@Override
				public boolean isAnonymousNode(IRI arg) {
					return false;
				}
				
			}, new OWLOntologyLoaderConfiguration());
			
			OWLAPIHandler converter = new SesameRDFUtils.OWLAPIHandler(consumer);
			
			atomQuery.evaluate(converter);
			atom.setAxiomSet(atomOntology.getAxioms());
		} catch (RepositoryException e) {
			throw new ADPersistenceException(e);
		} catch (MalformedQueryException e) {
			throw new RuntimeException(e);
		} catch (QueryEvaluationException e) {
			throw new ADPersistenceException(e);
		} catch (RDFHandlerException e) {
			throw new ADPersistenceException(e);
		} catch (OWLOntologyCreationException e) {
			throw new ADPersistenceException(e);
		}
	}

	@Override
	protected Reader getOntologySourceReader(PersistentAtom atom) throws ADPersistenceException {
		return null;
	} */


	@Override
	public void close() {
		//Don't close the connection because the AD could be using it
	}
}